#include "engine.h"

#include <fstream>

#include "NvInferRuntimeCommon.h"
#include "nvlogger.h"

namespace weilin {

bool Engine::build_engine(const std::string& file) {
  std::ifstream filein(file, std::ios::binary);
  filein.seekg(0, filein.end);
  size_t size = filein.tellg();
  filein.seekg(0, filein.beg);
  char* model_binary = new char[size];
  filein.read(model_binary, size);
  if (filein.bad()) {
    spdlog::error("Reading file {} error", file);
    return false;
  }
  auto res = this->build_engine(model_binary, size);
  delete[] model_binary;
  return res;
}

bool Engine::build_engine(const std::shared_ptr<nvinfer1::IHostMemory>& model) {
  return build_engine(static_cast<char*>(model->data()), model->size());
}

bool Engine::build_engine(const char* data, const std::size_t size) {
  std::unique_ptr<nvinfer1::IRuntime> runtime{
      nvinfer1::createInferRuntime(nv_logger)};
  if (runtime == nullptr) {
    spdlog::error("Create infer runtime error");
    return false;
  }

  mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(
      runtime->deserializeCudaEngine(data, size));
  perpare();
  return true;
}

void Engine::perpare() {
  mContext = std::unique_ptr<nvinfer1::IExecutionContext>(
      mEngine->createExecutionContext());
  mBuffers = std::shared_ptr<TensorBuffer>(new TensorBuffer(mEngine));
  isAvaliable = true;
}

std::shared_ptr<TensorBuffer>& Engine::get_buffer() {
  return mBuffers;
}

void Engine::execute() {
  get_buffer()->copy_input_to_device();
  mContext->executeV2(get_buffer()->get_device_bindings().data());
  get_buffer()->copy_output_to_host();
}

void Engine::execute_async() {
  // TODO: 使用DMA异步传输数据
  cudaStream_t stream;
  if (cudaStreamCreate(&stream) != cudaSuccess) {
    spdlog::error("Could not create cuda stream.");
    return;
  }

  get_buffer()->copy_input_to_device_async(stream);
  mContext->enqueue(1, get_buffer()->get_device_bindings().data(), stream,
                    nullptr);
  get_buffer()->copy_output_to_host_async(stream);
  cudaStreamSynchronize(stream);
  cudaStreamDestroy(stream);
}

}  // namespace weilin
